import numpy as np
import cv2 as cv
import imageio

import torch
import torch.nn.functional as F

from network.unet import Unet
from network.deeplab import DeepLabV1, DeepLabV2, DeepLabV3, DeepLabV3Plus
from network.danet import DANet, get_danet
from utils import *


if __name__ == '__main__':
    net = DeepLabV3Plus(nChannels=3, nClasses=1, nBlocks=[3, 4, 23, 3], atrousRates=[6, 12, 18, 24],
                        multiGrids=[1, 2, 4], outputStride=8)
    netFile = "./DeepLabV3Plus-best.pkl"
    data = imageio.imread("./wind1_37_4.jpg")
    data = np.swapaxes(data, 1, 2)
    data = np.swapaxes(data, 0, 1)
    data = np.expand_dims(data, axis=0) / 255.0

    device = torch.device("cpu")
    net.load_state_dict(torch.load(netFile, map_location=torch.device('cpu')))
    data = torch.from_numpy(data).float().to(device)
    output = net(data)
    # print(output)